package com.stars.easyms.datasource.mybatis;

import com.stars.easyms.datasource.EasyMsDataSource;
import com.stars.easyms.datasource.EasyMsMasterSlaveDataSource;
import com.stars.easyms.datasource.EasyMsMultiDataSource;
import com.stars.easyms.datasource.EasyMsQueryExecutor;
import com.stars.easyms.datasource.bean.PageCountDelayHandleBean;
import com.stars.easyms.datasource.holder.EasyMsMasterSlaveDataSourceHolder;
import com.stars.easyms.datasource.holder.EasyMsMybatisPlusHolder;
import com.stars.easyms.datasource.holder.EasyMsSynchronizationManager;
import com.stars.easyms.base.util.ReflectUtil;
import com.stars.easyms.datasource.transaction.EasyMsTransactionStatusHolder;
import com.stars.easyms.datasource.transaction.EasyMsTransactionSynchronizationManager;
import lombok.extern.slf4j.Slf4j;
import org.apache.ibatis.exceptions.PersistenceException;
import org.apache.ibatis.session.ExecutorType;
import org.apache.ibatis.session.SqlSession;
import org.apache.ibatis.session.SqlSessionFactory;
import org.mybatis.spring.SqlSessionTemplate;
import org.mybatis.spring.SqlSessionUtils;
import org.springframework.dao.DataAccessException;
import org.springframework.dao.TransientDataAccessResourceException;
import org.springframework.dao.support.PersistenceExceptionTranslator;
import org.springframework.jdbc.datasource.DataSourceUtils;
import org.springframework.lang.NonNull;
import org.springframework.lang.Nullable;
import org.springframework.transaction.support.TransactionSynchronizationAdapter;

import javax.sql.DataSource;
import java.lang.reflect.InvocationHandler;
import java.lang.reflect.Method;
import java.lang.reflect.Proxy;
import java.util.List;

import static org.apache.ibatis.reflection.ExceptionUtil.unwrapThrowable;

/**
 * <p>className: EasyMsSqlSessionTemplate</p>
 * <p>description: 重写SqlSessionTemplate类</p>
 *
 * @author guoguifang
 * @date 2019-09-20 11:01
 * @since 1.3.2
 */
@Slf4j
public class EasyMsSqlSessionTemplate extends SqlSessionTemplate {

    private final EasyMsMultiDataSource easyMsMultiDataSource;

    private final EasyMsQueryExecutor easyMsQueryExecutor = EasyMsQueryExecutor.getInstance();

    private final ExecutorType executorType;

    public EasyMsSqlSessionTemplate(EasyMsSqlSessionFactory sqlSessionFactory) {
        super(sqlSessionFactory, sqlSessionFactory.getConfiguration().getDefaultExecutorType(),
                new EasyMsMyBatisExceptionTranslator(
                        sqlSessionFactory.getConfiguration().getEnvironment().getDataSource(), true));
        this.executorType = sqlSessionFactory.getConfiguration().getDefaultExecutorType();
        DataSource dataSource = sqlSessionFactory.getConfiguration().getEnvironment().getDataSource();
        this.easyMsMultiDataSource = dataSource instanceof EasyMsMultiDataSource ? (EasyMsMultiDataSource) dataSource : null;
        try {
            ReflectUtil.writeFinalField(this, "sqlSessionProxy",
                    Proxy.newProxyInstance(SqlSessionFactory.class.getClassLoader(),
                            new Class[]{SqlSession.class},
                            new EasyMsSqlSessionInterceptor()));
        } catch (IllegalAccessException | NoSuchFieldException e) {
            log.error("There is no field 'sqlSessionProxy' in the class 'org.mybatis.spring.SqlSessionTemplate'!");
            System.exit(1);
        }
        sqlSessionFactory.setSqlSessionTemplate(this);
    }

    private class EasyMsSqlSessionInterceptor implements InvocationHandler {

        /**
         * 位于EasyMsRepositoryInterceptor、SpecifyDataSourceInterceptor之后，EasyMsQueryInterceptor、EasyMsUpdateInterceptor、EasyMsPageInterceptor之前，
         * 在此处分为两种情况：
         * 1.有事务：事务一般用于增删改，因此在该方法处强制设置为主数据源，不在后续设置
         * 2.无事务：因为在EasyMsUpdateInterceptor会强制将增删改置为主数据源，因此此处无法确定具体的数据源，需要在EasyMsQueryInterceptor、EasyMsUpdateInterceptor里具体得到数据源
         */
        @Override
        public Object invoke(Object proxy, Method method, Object[] args) throws Throwable {
            SqlSessionFactory sqlSessionFactory = getSqlSessionFactory();
            PersistenceExceptionTranslator exceptionTranslator = getPersistenceExceptionTranslator();

            // 判断是否是EasyMs事务的
            boolean isEasyMsTransactional = EasyMsTransactionSynchronizationManager.isSynchronizationActive() && easyMsMultiDataSource != null;

            Object result;
            EasyMsSqlSessionHolder sqlSessionHolder = null;
            SqlSession sqlSession = null;
            boolean isSetEasyMsMasterSlaveDataSource = false;

            // 进入同步管理状态
            EasyMsSynchronizationManager.enterSynchronization();
            try {
                // 如果是有事务的，则强制走主数据源
                if (isEasyMsTransactional) {
                    // 1.先判断主从数据源，先从线程本地变量中取，如果不存在则取默认主从数据源且需要再存放入线程本地变量中因为后续会用到
                    EasyMsMasterSlaveDataSource easyMsMasterSlaveDataSource = EasyMsMasterSlaveDataSourceHolder.getMasterSlaveDataSource();
                    if (easyMsMasterSlaveDataSource == null) {
                        easyMsMasterSlaveDataSource = easyMsMultiDataSource.getDefaultDataSource();
                        EasyMsMasterSlaveDataSourceHolder.putMasterSlaveDataSource(easyMsMasterSlaveDataSource);
                        isSetEasyMsMasterSlaveDataSource = true;
                    }

                    // 2.获取普通数据源，这里强制设置为主数据源，但是还是遵循负载均衡策略
                    EasyMsDataSource easyMsDataSource = EasyMsTransactionSynchronizationManager.getFixedMasterEasyMsDataSource(easyMsMasterSlaveDataSource);
                    EasyMsSynchronizationManager.setEasyMsDataSource(easyMsDataSource);

                    // 3.获取SqlSession，如果是事务的由于sqlSession是非线程安全的，因此每个线程每个普通数据源设置一个
                    sqlSessionHolder = getSqlSessionHolderInTransaction(
                            EasyMsTransactionSynchronizationManager.newTransactionStatusHolder(easyMsDataSource),
                            sqlSessionFactory, getExecutorType(), exceptionTranslator);
                    sqlSession = sqlSessionHolder.getSqlSession(sqlSessionFactory);
                } else {
                    sqlSession = SqlSessionUtils.getSqlSession(sqlSessionFactory, getExecutorType(), exceptionTranslator);
                }

                result = method.invoke(sqlSession, args);
                if (!isEasyMsTransactional) {
                    sqlSession.commit(true);
                }
                return handlePageCountDelay(EasyMsSynchronizationManager.getPageCountDelayHandleBean(), result);
            } catch (Throwable t) {
                Throwable unwrapped = unwrapThrowable(t);
                if (exceptionTranslator != null && unwrapped instanceof PersistenceException) {
                    if (sqlSession != null) {
                        if (isEasyMsTransactional) {
                            if (log.isDebugEnabled()) {
                                log.debug("Releasing transactional SqlSession [{}]", sqlSession);
                            }
                            sqlSessionHolder.released();
                        } else {
                            SqlSessionUtils.closeSqlSession(sqlSession, sqlSessionFactory);
                        }
                        sqlSession = null;
                    }
                    Throwable translated = exceptionTranslator.translateExceptionIfPossible((PersistenceException) unwrapped);
                    if (translated != null) {
                        unwrapped = translated;
                    }
                }
                throw unwrapped;
            } finally {
                if (sqlSession != null) {
                    if (isEasyMsTransactional) {
                        if (log.isDebugEnabled()) {
                            log.debug("Releasing transactional SqlSession [{}]", sqlSession);
                        }
                        if (sqlSessionHolder != null) {
                            sqlSessionHolder.released();
                        }
                    } else {
                        SqlSessionUtils.closeSqlSession(sqlSession, sqlSessionFactory);
                    }
                }
                EasyMsSynchronizationManager.clearSynchronization();

                // 如果在本方法中设置了主从数据源则在本方法内remove
                if (isSetEasyMsMasterSlaveDataSource) {
                    EasyMsMasterSlaveDataSourceHolder.clearMasterSlaveDataSource();
                }
            }
        }

        private Object handlePageCountDelay(@Nullable PageCountDelayHandleBean bean, Object result) {
            if (bean != null) {
                try {
                    easyMsQueryExecutor.handlePageCount(bean);
                    if (result instanceof List) {
                        result = bean.getDialect().afterPage((List) result, bean.getParameter(), bean.getRowBounds());
                    }
                } finally {
                    bean.getDialect().afterAll();
                }
            }
            return result;
        }

        /**
         * 因为sqlSession是非线程安全的，因此每个线程绑定一个sqlSession
         */
        @NonNull
        private EasyMsSqlSessionHolder getSqlSessionHolderInTransaction(@NonNull EasyMsTransactionStatusHolder easyMsTransactionStatusHolder,
                                                                        @NonNull SqlSessionFactory sqlSessionFactory, ExecutorType executorType,
                                                                        PersistenceExceptionTranslator exceptionTranslator) {
            EasyMsSqlSessionHolder sqlSessionHolder = easyMsTransactionStatusHolder.getSqlSessionHolder();

            // 如果线程本地变量中已经存在holder并且是处于事务中则正在使用次数加1，并返回sqlSession，允许用户自定义sqlSessionFactory因此需要将sqlSessionFactory与sqlSession绑定
            if (sqlSessionHolder != null && sqlSessionHolder.isSynchronizedWithTransaction()
                    && sqlSessionHolder.getSqlSession(sqlSessionFactory) != null) {
                if (sqlSessionHolder.getExecutorType() != executorType) {
                    throw new TransientDataAccessResourceException("Cannot change the ExecutorType when there is an existing transaction");
                }
                sqlSessionHolder.requested();
                if (log.isDebugEnabled()) {
                    log.debug("Fetched SqlSession [{}] from current transaction", sqlSessionHolder.getSqlSession(sqlSessionFactory));
                }
                return sqlSessionHolder;
            }

            // 如果SqlSessionHolder不存在或者SqlSessionHolder中没有保存SqlSession则创建一个新的SqlSession并保存
            if (log.isDebugEnabled()) {
                log.debug("Creating a new SqlSession");
            }
            SqlSession sqlSession = sqlSessionFactory.openSession(executorType);
            if (log.isDebugEnabled()) {
                log.debug("Registering transaction synchronization for SqlSession [{}]", sqlSession);
            }

            // 注册SqlSession的holder到线程本地变量中
            if (sqlSessionHolder != null && sqlSessionHolder.isSynchronizedWithTransaction()) {
                sqlSessionHolder.setSqlSession(sqlSessionFactory, sqlSession);
            } else {
                sqlSessionHolder = new EasyMsSqlSessionHolder(sqlSessionFactory, sqlSession, executorType, exceptionTranslator);
                easyMsTransactionStatusHolder.setSqlSessionHolder(sqlSessionHolder);
                easyMsTransactionStatusHolder.addSynchronizations(new EasyMsSqlSessionSynchronization(sqlSessionHolder));
                sqlSessionHolder.setSynchronizedWithTransaction(true);
            }
            sqlSessionHolder.requested();
            return sqlSessionHolder;
        }

        private final class EasyMsSqlSessionSynchronization extends TransactionSynchronizationAdapter {

            private final EasyMsSqlSessionHolder holder;

            private boolean holderActive = true;

            private EasyMsSqlSessionSynchronization(@NonNull EasyMsSqlSessionHolder holder) {
                this.holder = holder;
            }

            @Override
            public int getOrder() {
                return DataSourceUtils.CONNECTION_SYNCHRONIZATION_ORDER - 1;
            }

            @Override
            public void suspend() {
                if (this.holderActive && log.isDebugEnabled()) {
                    for (SqlSession sqlSession : this.holder.getSqlSessionSet()) {
                        log.debug("Transaction synchronization suspending SqlSession [{}]", sqlSession);
                    }
                }
            }

            @Override
            public void resume() {
                if (this.holderActive && log.isDebugEnabled()) {
                    for (SqlSession sqlSession : this.holder.getSqlSessionSet()) {
                        log.debug("Transaction synchronization resuming SqlSession [{}]", sqlSession);
                    }
                }
            }

            @Override
            public void beforeCommit(boolean readOnly) {
                if (EasyMsTransactionSynchronizationManager.isActualTransactionActive()) {
                    try {
                        for (SqlSession sqlSession : this.holder.getSqlSessionSet()) {
                            if (log.isDebugEnabled()) {
                                log.debug("Transaction synchronization committing SqlSession [{}]", sqlSession);
                            }
                            sqlSession.commit();
                        }
                    } catch (PersistenceException p) {
                        if (this.holder.getPersistenceExceptionTranslator() != null) {
                            DataAccessException translated = this.holder
                                    .getPersistenceExceptionTranslator()
                                    .translateExceptionIfPossible(p);
                            if (translated != null) {
                                throw translated;
                            }
                        }
                        throw p;
                    }
                }
            }

            @Override
            public void beforeCompletion() {
                if (!this.holder.isOpen()) {
                    close();
                }
            }

            @Override
            public void afterCompletion(int status) {
                if (this.holderActive) {
                    close();
                }
                this.holder.reset();
            }

            private void close() {
                this.holderActive = false;
                for (SqlSession sqlSession : this.holder.getSqlSessionSet()) {
                    if (log.isDebugEnabled()) {
                        log.debug("Transaction synchronization deregistering SqlSession [{}]", sqlSession);
                    }
                    if (log.isDebugEnabled()) {
                        log.debug("Transaction synchronization closing SqlSession [{}]", sqlSession);
                    }
                    sqlSession.close();
                }
            }
        }
    }

    @Override
    public ExecutorType getExecutorType() {
        ExecutorType localExecutorType = EasyMsMybatisPlusHolder.getExecutorType();
        if (localExecutorType != null) {
            return localExecutorType;
        }
        return this.executorType;
    }

    @Override
    public void commit() {
        if (EasyMsMybatisPlusHolder.getExecutorType() == null) {
            super.commit();
        }
    }

    @Override
    public void commit(boolean force) {
        if (EasyMsMybatisPlusHolder.getExecutorType() == null) {
            super.commit(force);
        }
    }

    @Override
    public void rollback() {
        if (EasyMsMybatisPlusHolder.getExecutorType() == null) {
            super.rollback();
        }
    }

    @Override
    public void rollback(boolean force) {
        if (EasyMsMybatisPlusHolder.getExecutorType() == null) {
            super.rollback(force);
        }
    }

    @Override
    public void close() {
        if (EasyMsMybatisPlusHolder.getExecutorType() == null) {
            super.close();
        }
    }
}